package com.treasure.web.service.impl;

import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.google.common.collect.Lists;
import com.treasure.common.enums.TrainingStatEnum;
import com.treasure.common.exception.ParameterException;
import com.treasure.common.exception.ServiceException;
import com.treasure.common.model.SessionUser;
import com.treasure.common.model.SessionUserHolder;
import com.treasure.common.types.Corp;
import com.treasure.domain.domain.entity.document.UploadFileEntity;
import com.treasure.common.model.UploadDocumentRequest;
import com.treasure.infrastructure.model.po.CorpUserInfo;
import com.treasure.infrastructure.model.po.KnowledgeBase;
import com.treasure.infrastructure.model.po.KnowledgeDocument;
import com.treasure.infrastructure.model.po.KnowledgeDocumentVector;
import com.treasure.infrastructure.repository.dao.mapper.CorpUserInfoMapper;
import com.treasure.infrastructure.repository.dao.mapper.KnowledgeBaseMapper;
import com.treasure.infrastructure.repository.dao.mapper.KnowledgeDocumentMapper;
import com.treasure.infrastructure.repository.dao.mapper.KnowledgeDocumentVectorMapper;
import com.treasure.web.command.event.DocumentUploadEvent;
import com.treasure.web.command.event.EventPublisher;
import com.treasure.web.command.request.KnowledgeDocumentOperateRequest;
import com.treasure.web.command.request.KnowledgeDocumentRequest;
import com.treasure.web.command.response.DocumentDownloadResponse;
import com.treasure.web.command.response.KnowledgeDocumentResponse;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.treasure.web.command.response.KnowledgeDocumentVectorResponse;
import com.treasure.web.command.response.KnowledgeRepositoryResponse;
import com.treasure.web.mapstruct.KnowledgeDocumentAssembler;
import com.treasure.web.mapstruct.KnowledgeDocumentVectorAssembler;
import com.treasure.web.mapstruct.KnowledgeRepositoryAssembler;
import com.treasure.web.service.IKnowledgeBaseService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.pdfbox.io.IOUtils;
import org.jetbrains.annotations.NotNull;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.web.multipart.MultipartFile;

import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.time.LocalDateTime;
import java.util.*;
import java.util.stream.Collectors;

/**
 * <p>
 * 服务实现类
 * </p>
 *
 * @author sunhongmin
 * @since 2023-08-31
 */
@Service
@RequiredArgsConstructor
@Slf4j
public class KnowledgeBaseServiceImpl extends ServiceImpl<KnowledgeBaseMapper, KnowledgeBase> implements IKnowledgeBaseService {

    private final KnowledgeBaseMapper knowledgeBaseMapper;
    private final CorpUserInfoMapper corpUserInfoMapper;
    private final KnowledgeDocumentMapper knowledgeDocumentMapper;
    private final KnowledgeDocumentVectorMapper knowledgeDocumentVectorMapper;

    @Override
    public List<KnowledgeRepositoryResponse> getKnowledgeRepositoryList(Corp corpId) {

        QueryWrapper<KnowledgeBase> knowledgeBaseQueryWrapper = new QueryWrapper<>();
        knowledgeBaseQueryWrapper.eq("corp_id", corpId.getValue());

        List<KnowledgeBase> knowledgeBaseList = knowledgeBaseMapper.selectList(knowledgeBaseQueryWrapper);

        if (CollectionUtils.isEmpty(knowledgeBaseList)) {
            return Collections.emptyList();
        }

        ArrayList<String> userIdList = knowledgeBaseList.stream().map(KnowledgeBase::getCreatorId).distinct().collect(Collectors.toCollection(ArrayList::new));
        Map<String, String> userIdNameMap = getUserIdToNameMap(corpId, userIdList);

        List<KnowledgeRepositoryResponse> knowledgeRepositoryResponseList = KnowledgeRepositoryAssembler.INSTANCE.convertToKnowledgeRepositoryResponseList(knowledgeBaseList);
        for (KnowledgeRepositoryResponse knowledgeRepositoryResponse : knowledgeRepositoryResponseList) {
            knowledgeRepositoryResponse.setCreatorName(userIdNameMap.get(knowledgeRepositoryResponse.getCreatorId()));
        }

        return knowledgeRepositoryResponseList;
    }

    @NotNull
    private Map<String, String> getUserIdToNameMap(Corp corpId, ArrayList<String> userIdList) {
        QueryWrapper<CorpUserInfo> corpUserInfoQueryWrapper = new QueryWrapper<CorpUserInfo>();
        corpUserInfoQueryWrapper.eq("corp_id", corpId.getValue())
                .in("user_id", userIdList);
        if (CollectionUtils.isEmpty(userIdList)) {
            return Collections.emptyMap();
        }
        List<CorpUserInfo> corpUserInfoList = corpUserInfoMapper.selectList(corpUserInfoQueryWrapper);
        return Optional.ofNullable(corpUserInfoList).orElse(Collections.emptyList()).stream().collect(Collectors.toMap(CorpUserInfo::getUserId, CorpUserInfo::getUsername));
    }

    @Override
    public List<KnowledgeDocumentResponse> getDocumentList(Corp corpId, KnowledgeDocumentRequest knowledgeDocumentRequest) {
        knowledgeDocumentRequest.valiadate();

        QueryWrapper<KnowledgeDocument> knowledgeDocumentQueryWrapper = new QueryWrapper<>();
        knowledgeDocumentQueryWrapper.eq("corp_id", corpId.getValue())
                .eq("knowledge_id", knowledgeDocumentRequest.getKnowledgeId());

        List<KnowledgeDocument> knowledgeDocumentList = knowledgeDocumentMapper.selectList(knowledgeDocumentQueryWrapper);
        if (CollectionUtils.isEmpty(knowledgeDocumentList)) {
            return Collections.emptyList();
        }

        ArrayList<String> userIdList = knowledgeDocumentList.stream().map(KnowledgeDocument::getCreatorId).distinct().collect(Collectors.toCollection(ArrayList::new));
        Map<String, String> userIdNameMap = getUserIdToNameMap(corpId, userIdList);

        List<KnowledgeDocumentResponse> knowledgeDocumentResponsesList = KnowledgeRepositoryAssembler.INSTANCE.convertToKnowledgeDocumentResponseList(knowledgeDocumentList);
        for (KnowledgeDocumentResponse knowledgeDocumentResponse : knowledgeDocumentResponsesList) {
            knowledgeDocumentResponse.setCreatorName(userIdNameMap.get(knowledgeDocumentResponse.getCreatorId()));
        }

        return knowledgeDocumentResponsesList;
    }

    @Override
    @Transactional
    public KnowledgeDocumentResponse uploadDocument(MultipartFile file, UploadDocumentRequest uploadDocumentRequest) throws IOException {
        UploadFileEntity uploadFile = new UploadFileEntity(file);
        log.info("uploadDocument fileName:{},fileSize:{}", uploadFile.getFileName().getValue(), uploadFile.getFileSize());

        // 知识库校验
        KnowledgeBase knowledgeBase = getKnowledgeBase(uploadDocumentRequest);
        if (Objects.isNull(knowledgeBase)) {
            log.info("uploadDocument 知识库不存在 id={}", uploadDocumentRequest.getKnowledgeId());
            throw new ServiceException("指定的知识库不存在");
        }

        // 名称获取相同文档，替换更新
        KnowledgeDocument knowledgeDocument = checkAndGetDocument(uploadDocumentRequest, uploadFile);
        uploadDocumentRequest.setDocumentId(knowledgeDocument.getId());
        uploadDocumentRequest.setDocumentName(knowledgeDocument.getDocumentName());

        // 推送文档训练事件
        DocumentUploadEvent documentUploadEvent = new DocumentUploadEvent(this, uploadDocumentRequest, uploadFile);
        EventPublisher.applicationEventPublisher.publishEvent(documentUploadEvent);

        // 转换返回结果
        KnowledgeDocumentResponse knowledgeDocumentResponse = KnowledgeDocumentAssembler.INSTANCE.convertToKnowledgeDocumentResponse(knowledgeDocument);
        knowledgeDocumentResponse.setCreatorName(uploadDocumentRequest.getUserName());
        return knowledgeDocumentResponse;
    }

    private KnowledgeBase getKnowledgeBase(UploadDocumentRequest uploadDocumentRequest) {
        SessionUser sessionUser = SessionUserHolder.getAndCheckSessionUser();
        QueryWrapper<KnowledgeBase> knowledgeQuery = new QueryWrapper<>();
        knowledgeQuery.lambda().eq(KnowledgeBase::getCorpId, sessionUser.getCorpId())
                .eq(KnowledgeBase::getId, uploadDocumentRequest.getKnowledgeId());
        return knowledgeBaseMapper.selectOne(knowledgeQuery);
    }

    @NotNull
    private KnowledgeDocument checkAndGetDocument(UploadDocumentRequest uploadDocumentRequest, UploadFileEntity uploadFile) {
        SessionUser sessionUser = SessionUserHolder.getAndCheckSessionUser();
        QueryWrapper<KnowledgeDocument> documentQuery = new QueryWrapper<>();
        documentQuery.lambda().eq(KnowledgeDocument::getCorpId, sessionUser.getCorpId())
                .eq(KnowledgeDocument::getDocumentName, uploadFile.getFileName().getValue())
                .eq(KnowledgeDocument::getKnowledgeId, uploadDocumentRequest.getKnowledgeId())
                .eq(KnowledgeDocument::getDocumentSuffix, uploadFile.getFileName().getFileSuffix());
        KnowledgeDocument documentExists = knowledgeDocumentMapper.selectOne(documentQuery);
        if (Objects.nonNull(documentExists)) {
            // 训练状态-初始化
            documentExists.setTrainingStat(TrainingStatEnum.IN_TRAINING.getValue());
            knowledgeDocumentMapper.updateById(documentExists);
            return documentExists;
        } else {
            KnowledgeDocument knowledgeDocument = getKnowledgeDocument(uploadDocumentRequest, uploadFile.getFileSize(), uploadFile.getFileName().getValue(), uploadFile.getFileName().getFileSuffix());
            // 训练状态-初始化
            knowledgeDocument.setTrainingStat(TrainingStatEnum.IN_TRAINING.getValue());
            knowledgeDocumentMapper.insert(knowledgeDocument);
            return knowledgeDocument;
        }
    }


    @Override
    public void deleteDocument(KnowledgeDocumentOperateRequest knowledgeDocumentDeleteRequest) {
        SessionUser sessionUser = SessionUserHolder.getAndCheckSessionUser();
        KnowledgeDocument knowledgeDocument = getKnowledgeDocument(knowledgeDocumentDeleteRequest, sessionUser);
        if (Objects.isNull(knowledgeDocument)) {
            log.error("deleteDocument not found document. documentId:{}", knowledgeDocumentDeleteRequest.getDocumentId());
            return;
        }
        knowledgeDocumentMapper.deleteById(knowledgeDocument.getId());

        QueryWrapper<KnowledgeDocumentVector> vectorQuery = new QueryWrapper<>();
        vectorQuery.lambda().select(Lists.newArrayList(KnowledgeDocumentVector::getId));
        vectorQuery.lambda().eq(KnowledgeDocumentVector::getCorpId, sessionUser.getCorpId())
                .eq(KnowledgeDocumentVector::getDocumentId, knowledgeDocument.getId());
        List<KnowledgeDocumentVector> knowledgeDocumentVectors = knowledgeDocumentVectorMapper.selectList(vectorQuery);
        if (CollectionUtils.isEmpty(knowledgeDocumentVectors)) {
            log.error("deleteDocument not found document vector. documentId:{}", knowledgeDocumentDeleteRequest.getDocumentId());
            return;
        }

        List<Long> vectorIds = knowledgeDocumentVectors.stream().map(KnowledgeDocumentVector::getId).collect(Collectors.toList());
        knowledgeDocumentVectorMapper.deleteBatchIds(vectorIds);
    }

    private KnowledgeDocument getKnowledgeDocument(KnowledgeDocumentOperateRequest knowledgeDocumentDeleteRequest, SessionUser sessionUser) {
        QueryWrapper<KnowledgeDocument> documentQuery = new QueryWrapper<>();
        documentQuery.eq("id", knowledgeDocumentDeleteRequest.getDocumentId())
                .eq("corp_id", sessionUser.getCorpId());
        return knowledgeDocumentMapper.selectOne(documentQuery);
    }

    @Override
    public List<KnowledgeDocumentVectorResponse> getDocumentVectors(KnowledgeDocumentOperateRequest knowledgeDocumentDeleteRequest) {
        SessionUser sessionUser = SessionUserHolder.getAndCheckSessionUser();
        KnowledgeDocument knowledgeDocument = getKnowledgeDocument(knowledgeDocumentDeleteRequest, sessionUser);
        if (Objects.isNull(knowledgeDocument)) {
            throw new ServiceException("deleteDocument not found document. documentId:" + knowledgeDocumentDeleteRequest.getDocumentId());
        }
        QueryWrapper<KnowledgeDocumentVector> vectorQuery = new QueryWrapper<>();
        vectorQuery.lambda().select(Lists.newArrayList(KnowledgeDocumentVector::getId, KnowledgeDocumentVector::getDocumentId, KnowledgeDocumentVector::getDocumentName, KnowledgeDocumentVector::getCorpId, KnowledgeDocumentVector::getPassages, KnowledgeDocumentVector::getCreateTime, KnowledgeDocumentVector::getCreatorId, KnowledgeDocumentVector::getHasDeleted));
        vectorQuery.lambda().eq(KnowledgeDocumentVector::getCorpId, sessionUser.getCorpId())
                .eq(KnowledgeDocumentVector::getDocumentId, knowledgeDocument.getId())
                .eq(KnowledgeDocumentVector::getHasDeleted, Boolean.FALSE);
        List<KnowledgeDocumentVector> knowledgeDocumentVectors = knowledgeDocumentVectorMapper.selectList(vectorQuery);
        List<KnowledgeDocumentVectorResponse> knowledgeDocumentVectorResponses = KnowledgeDocumentVectorAssembler.INSTANCE.convertResponseList(knowledgeDocumentVectors);
        if (CollectionUtils.isEmpty(knowledgeDocumentVectorResponses)) {
            for (KnowledgeDocumentVectorResponse knowledgeDocumentVectorResponse : knowledgeDocumentVectorResponses) {
                knowledgeDocumentVectorResponse.setKnowledgeId(knowledgeDocument.getKnowledgeId());
            }
        }
        return knowledgeDocumentVectorResponses;
    }

    @Override
    public DocumentDownloadResponse downloadDocument(KnowledgeDocumentOperateRequest knowledgeDocumentDeleteRequest) throws IOException {
        SessionUser sessionUser = SessionUserHolder.getAndCheckSessionUser();
        KnowledgeDocument knowledgeDocument = getKnowledgeDocument(knowledgeDocumentDeleteRequest, sessionUser);
        if (Objects.isNull(knowledgeDocument)) {
            log.error("deleteDocument not found document. documentId:{}", knowledgeDocumentDeleteRequest.getDocumentId());
            throw new ParameterException("未查询到文档信息");
        }

        String knowledgeFilePath = UploadFileEntity.getKnowledgeFilePath();
        String filePath = knowledgeFilePath + File.separator + sessionUser.getCorpId() + File.separator + knowledgeDocument.getDocumentName();
        return new DocumentDownloadResponse(knowledgeDocument.getDocumentName(), IOUtils.toByteArray(Files.newInputStream(Paths.get(filePath))));
    }

    @NotNull
    private static KnowledgeDocument getKnowledgeDocument(UploadDocumentRequest uploadDocumentRequest, long fileSize, String fileName, String fileSuffix) {
        KnowledgeDocument knowledgeDocument = new KnowledgeDocument();
        knowledgeDocument.setCorpId(uploadDocumentRequest.getCorpId());
        knowledgeDocument.setKnowledgeId(uploadDocumentRequest.getKnowledgeId());
        knowledgeDocument.setDocumentName(fileName);
        knowledgeDocument.setDocumentSuffix(fileSuffix);
        knowledgeDocument.setDocumentSize(Long.valueOf(fileSize).intValue());
        knowledgeDocument.setCreateTime(LocalDateTime.now());
        knowledgeDocument.setCreatorId(uploadDocumentRequest.getUserId());
        return knowledgeDocument;
    }
}
